package ipmanager

import (
	"fmt"
	"net"
	"sync"
)

type PerIPConnCounter struct {
	pool sync.Pool
	lock sync.Mutex
	m    map[uint32]int
}

func (cc *PerIPConnCounter) Register(ip uint32) int {
	cc.lock.Lock()
	if cc.m == nil {
		cc.m = make(map[uint32]int)
	}
	n := cc.m[ip] + 1
	cc.m[ip] = n
	cc.lock.Unlock()
	return n
}

func (cc *PerIPConnCounter) Unregister(ip uint32) {
	cc.lock.Lock()
	if cc.m == nil {
		cc.lock.Unlock()
		panic("BUG: perIPConnCounter.Register() wasn't called")
	}
	n := cc.m[ip] - 1
	if n < 0 {
		cc.lock.Unlock()
		panic(fmt.Sprintf("BUG: negative per-ip counter=%d for ip=%d", n, ip))
	}
	cc.m[ip] = n
	cc.lock.Unlock()
}

type PerIPConn struct {
	net.Conn

	ip               uint32
	perIPConnCounter *PerIPConnCounter
}

func acquirePerIPConn(conn net.Conn, ip uint32, counter *PerIPConnCounter) *PerIPConn {
	v := counter.pool.Get()
	if v == nil {
		v = &PerIPConn{
			perIPConnCounter: counter,
		}
	}
	c := v.(*PerIPConn)
	c.Conn = conn
	c.ip = ip
	return c
}

func releasePerIPConn(c *PerIPConn) {
	c.Conn = nil
	c.perIPConnCounter.pool.Put(c)
}

func (c *PerIPConn) Close() error {
	if c == nil {
		return nil
	}
	if c.Conn == nil {
		return nil
	}
	err := c.Conn.Close()
	c.perIPConnCounter.Unregister(c.ip)
	releasePerIPConn(c)
	return err
}
